import base64
import json

from captcha.models import CaptchaStore
from captcha.views import captcha_image
from django.db.models import Q
from django.http import HttpResponse
from rest_framework.views import APIView

from ai.serializers import MyTokenObtainPairSerializer
from rest_framework_simplejwt.views import TokenObtainPairView
from rest_framework.permissions import AllowAny
from django.contrib.auth import get_user_model
from django.contrib.auth.backends import ModelBackend

User = get_user_model()


class MyCustomBackend(ModelBackend):
    def authenticate(self, request, username=None, password=None, **kwargs):
        try:
            user = User.objects.get(Q(username=username) | Q(email=username))
            if user.check_password(password):
                return user
        except Exception as e:
            return None


class CaptchaAPIView(APIView):

    def get(self, request):
        hashkey = CaptchaStore.generate_key()
        try:
            id_ = CaptchaStore.objects.filter(hashkey=hashkey).first().id
            image = captcha_image(request, hashkey)
            image_base = 'data:image/png;base64,%s' % base64.b64encode(image.content).decode('utf=8')
            json_data = json.dumps({'id': id_, 'image': image_base})
            CaptchaStore.remove_expired()
        except:
            json_data = None
        return HttpResponse(json_data, content_type='application/json')


class MyObtainTokenPairView(TokenObtainPairView):
    permission_classes = (AllowAny,)
    serializer_class = MyTokenObtainPairSerializer
